from __future__ import division
from __future__ import print_function

import os.path as osp
import numpy as np
from easydict import EasyDict as edict

__C = edict()
cfg = __C

# Dataset name: flowers, birds
__C.DATASET_NAME = 'birds'  # 与下面的CONFIG_NAME作用类似，作为当前训练数据超参数等等的ID，对程序运行无实际影响。
__C.EMBEDDING_TYPE = 'cnn-rnn'
'''
出现在datesets.py中，是TextDataset类的__init__函数中的一个参数，默认为'cnn-rnn'，还有其他的比如'cnn-gru','skip-thought'，
其取值影响TextDataSet类在实例化时导入的文件地址。我感觉这个类型是由encoder来确定的。
'''
__C.CONFIG_NAME = ''  # 可以理解为给CONFIG文件的一个ID，其取值对程序运行无任何影响，只是一个记号。与DATASET_NAME和timestamp连用，见上。
__C.GPU_ID = '0'  # 顾名思义，分给各个gpu一个id。其数据类型为str，各个id之间以逗号隔开（后面要对此以逗号分词变为列表的，列表长度即为gpu数量）。
__C.CUDA = True  # 顾名思义。看你用不用cuda。
__C.WORKERS = 6
'''
用于数据加载的子进程数。0意味着数据将在主进程中加载。（默认值：0）
我不懂这个东西。
上面这个超参量是torch.utils.data.DataLoader函数的一个参数，dataloader是前者一个实例，
dataloader可以作为 GANTrainer中函数train 的第一个参数。
'''

__C.NET_G = ''
'''
它是一个地址，该地址下存放着一个文件，这个文件可以通过torch.load函数转换成一个字典，
然后再通过netG.load_state_dict函数把刚才那个字典重新映射成一个完整的netG。
https://pytorch.org/docs/stable/generated/torch.load.html?highlight=torch%20load#torch.load
https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict
'''
__C.NET_D = ''  # 同上
__C.STAGE1_G = ''  # 类似上面，用于将存好的STAGE1_G赋值给netG.STAGE1_G
__C.DATA_DIR = 'data'  # 一个目录，很难说清，建议直接看TextDataset的__init__函数。后跟split的值。
__C.VIS_COUNT = 64  # 需要保存的原图片的高度

__C.Z_DIM = 100  # 噪声noise的尺寸为：(batch_size,Z_DIM)
'''
我感觉其实这个超参数具体的取值在小范围内的变化不会影响结果。
'''
__C.IMSIZE = 64  # 图片尺寸
__C.STAGE = 1  # 顾名思义，第几个阶段。不管TRAIN_FLAG是多少，这个超参数都有用。

# Training options
__C.TRAIN = edict()  # easydict
__C.TRAIN.FLAG = True  # main函数中，决定当前是训练还是测试。
__C.TRAIN.BATCH_SIZE = 64  # 分到每个gpu上的一次数据load的batch_size，因此实际load数据的时候一次的数据量是该超参量与gpu_num的乘积。
__C.TRAIN.MAX_EPOCH = 600  # 训练最大epoch
__C.TRAIN.SNAPSHOT_INTERVAL = 50  # 顾名思义，快照保存间隔。每次快照记录netG，netD，epoch。
__C.TRAIN.PRETRAINED_MODEL = ''  # 见下
__C.TRAIN.PRETRAINED_EPOCH = 600
'''
上面两个超参量都没有出现在程序中，合理怀疑它在models文件夹下面。这应该是预训练好的模型。我尝试点击下载链接，但还是要写那个申请。吐了。
'''
__C.TRAIN.LR_DECAY_EPOCH = 600  # 学习率decay的epoch
__C.TRAIN.DISCRIMINATOR_LR = 2e-4  # 判别器D的学习率
__C.TRAIN.GENERATOR_LR = 2e-4  # 生成器G的学习率

__C.TRAIN.COEFF = edict()  # easydict
__C.TRAIN.COEFF.KL = 2.0  # 我认为是正则化系数。

# Modal options
__C.GAN = edict()  # easydict
__C.GAN.CONDITION_DIM = 128  # 条件增强的维度。
__C.GAN.DF_DIM = 64  # 判别器一开始的通道数初始化参数
__C.GAN.GF_DIM = 128  # 生成器一开始的通道数初始化参数
__C.GAN.R_NUM = 4  # 残差网络块堆叠系数

__C.TEXT = edict()  # easydict
__C.TEXT.DIMENSION = 1024  # embedding的维度


def _merge_a_into_b(a, b):
    """Merge config dictionary a into config dictionary b, clobbering the
    options in b whenever they are also specified in a.
    """
    if type(a) is not edict:
        return

    for k, v in a.items():
        # a must specify keys that are in b
        if k not in b:
            raise KeyError('{} is not a valid config key'.format(k))

        # the types must match, too
        old_type = type(b[k])
        if old_type is not type(v):
            if isinstance(b[k], np.ndarray):
                v = np.array(v, dtype=b[k].dtype)
            else:
                raise ValueError(('Type mismatch ({} vs. {}) '
                                  'for config key: {}').format(type(b[k]),
                                                               type(v), k))

        # recursively merge dicts
        if type(v) is edict:
            try:
                _merge_a_into_b(a[k], b[k])
            except:
                print('Error under config key: {}'.format(k))
                raise
        else:
            b[k] = v


def cfg_from_file(filename):
    """Load a config file and merge it into the default options."""
    import yaml
    with open(filename, 'r') as f:
        yaml_cfg = edict(yaml.load(f, Loader=yaml.FullLoader))

    _merge_a_into_b(yaml_cfg, __C)
